import torch
from torch.nn.functional import softmax

x = [
    [1,0,1,0],
    [0,2,0,2],
    [1,1,1,1]
]
x = torch.tensor(x,dtype=torch.float32)
#初始化参数矩阵
w_key = [
    [0,0,1],
    [1,1,0],
    [0,1,0],
    [1,1,0]
]
w_query =[
[1,0,1],
    [1,0,0],
    [0,0,1],
    [0,1,1]
]
w_value = [
    [0,2,0],
    [0,3,0],
    [1,0,3],
    [1,1,0]
]
w_key = torch.tensor(w_key,dtype=torch.float32)
w_query = torch.tensor(w_query,dtype=torch.float32)
w_value = torch.tensor(w_value,dtype=torch.float32)

keys = x @ w_key
querys = x @ w_query
values = x @ w_value
# print(keys)

attn_scores = querys @ keys.T
print(attn_scores)

attn_scores_softmax = softmax(attn_scores,dim=-1)
print(attn_scores_softmax)